import pathlib

import matplotlib.pyplot
import numpy

from utilities.checkpoint_manager import CheckpointManager


def main(checkpoint_manager: CheckpointManager,
         plots: tuple[tuple[int], ...] = ((0, 500, 1000), (1500, 2000, 2500))):
    total_row = len(plots)
    total_column = len(plots[0])
    for row_index, row in enumerate(plots):
        for column_index, t in enumerate(row, start=1):
            matplotlib.pyplot.subplot(
                total_row, total_column,
                row_index * total_column + column_index)

            checkpoint_path = checkpoint_manager.existed()[t].with_suffix(".npy")
            checkpoint = numpy.load(checkpoint_path)

            t = t * 0.0001
            x = checkpoint[:, 0]
            y = numpy.sin(numpy.pi * x) * numpy.exp(-numpy.pi ** 2 * t)
            matplotlib.pyplot.plot(x, y, lw=5, alpha=0.2, label="Expected")

            y = checkpoint[:, 1]
            matplotlib.pyplot.plot(x, y, label="Actual")

            matplotlib.pyplot.legend(loc="upper right")
            # noinspection SpellCheckingInspection
            matplotlib.pyplot.title(rf"$t={t:.2f}$")
            matplotlib.pyplot.ylim(0, 1)

    matplotlib.pyplot.subplots_adjust(
        left=0.038, bottom=0.05, right=0.964, top=0.955,
        wspace=0.217, hspace=0.21)
    matplotlib.pyplot.show()


if __name__ == "__main__":
    main(CheckpointManager(pathlib.Path("./outputs")))
